from typing import Optional

from .pdf_parser import PDFParser
from .writer import Writer
from .translator_chain import TranslatorChain

from utils import LOG
from book import ContentType

class PDFTranslator:
    def __init__(self, model_type: str, model: str):
        self.book = None
        self.model_type = model_type
        self.translator_chain = TranslatorChain(model_type, model)
        self.pdf_parser = PDFParser()
        self.writer = Writer()

    def translate_pdf(self, input_file_path: str, output_file_format: str ='markdown', output_file_path: str = None, target_language: str = 'Chinese', page_num: Optional[int] = None):
        self.book = self.pdf_parser.parse_pdf(input_file_path, page_num)
        LOG.info(f'{input_file_path} 解析完成')

        LOG.info(f'开始翻译 {input_file_path}')
        for page_idx, page in enumerate(self.book.pages):
            for content_idx, content in enumerate(page.contents):
                LOG.info(content.original)

                if content.content_type == ContentType.IMAGE:
                    content.status = True
                    continue

                translation, status = self.translator_chain.invoke(content.original, target_language)
                LOG.info(translation)

                if self.model_type == "Tongyi":
                    # Tongyi 返回的是字符串，而不是AIMessage
                    self.book.pages[page_idx].contents[content_idx].set_translation(translation, status)
                else:
                    self.book.pages[page_idx].contents[content_idx].set_translation(translation.content, status)
                    
        LOG.info(f'翻译完成 {input_file_path}')

        LOG.info(f'开始保存 {output_file_path}')
        self.writer.save_translated_book(self.book, output_file_path, output_file_format)
        LOG.info(f'保存完成 {output_file_path}')